from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# 加载模型和分词器
model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# 创建 pipeline
generator = pipeline('text-generation', model=model, tokenizer=tokenizer)

# 测试
story = generator("Once upon a time in a land far away", max_length=1000, num_return_sequences=1)
print(story[0]['generated_text'])
